import numpy as np
import matplotlib.pyplot as plt
from sko.GA import GA
import pandas as pd

# 首先学习一下官方文档

'''
min f(x1, x2, x3) = x1^2 + x2^2 + x3^2
s.t.
    x1*x2 >= 1
    x1*x2 <= 5
    x2 + x3 = 1
    0 <= x1, x2, x3 <= 5
'''


def obj_func(p):
    print(p)
    x1, x2, x3 = p
    return x1 ** 2 + x2 ** 2 + x3 ** 2

if __name__ == '__main__':

    constraint_eq = [
        # 等式约束条件：
        lambda x: 15 - x[1] - x[2]
    ]

    constraint_ueq = [
        # 不等式约束条件：
        lambda x: 22 - x[0] * x[1],
        lambda x: x[0] * x[1] - 60
        # 无需使用不等式表示，直接使用表达式即可，但该表达式需要小于等于0，列如第一个约束条件5>=x1*x2,则改写成x1*x2-5：
    ]

    ga = GA(func=obj_func, n_dim=3, size_pop=50, max_iter=100, lb=[0, 4, 8], ub=[3, 7, 10], constraint_eq=constraint_eq,
            constraint_ueq=constraint_ueq, precision=[1, 1, 1])
    # func：需要优化的目标，n_dim:优化参数个数，lb、ub参数的上下界,constraint_eq、constraint_ueq：约束条件，
    # precision：可以使用整数或浮点数，若使用整数例如1则自变量会启用整数规划模式，浮点数同理，但是优化结果得到的整数是带有小数点的（1.0000），在有些情况下比如random_state他的值必须是个整数，此时需要将结果转成整型
    best_x, best_y = ga.run()
    print('best_x:', best_x, '\n', 'best_y:', best_y)

    Y_history = pd.DataFrame(ga.all_history_Y)
    fig, ax = plt.subplots(2, 1)
    ax[0].plot(Y_history.index, Y_history.values, '.', color='red')
    Y_history.min(axis=1).cummin().plot(kind='line')
    plt.show()